
from PIL import Image, ImageDraw, ImageFont

# import bs4
import cv2
import glob
import logging
import matplotlib.pyplot as plt
import os

import helper


logger = logging.getLogger(__name__)


class Main():

    def __init__(self, args):
        self.args = args

    def start(self):
        font = ImageFont.truetype('arial.ttf', 16)

        class_names = []

        if self.args.out_dir == self.args.img_dir:
            raise ValueError('out_dir and img_dir must be different')
        os.makedirs(self.args.out_dir, exist_ok=True)

        for pic_fullname in glob.glob(os.path.join(self.args.img_dir, f'*{self.args.img_ext}')):
            basename = os.path.basename(pic_fullname)

            img = Image.open(pic_fullname).convert('RGBA')
            iw, ih = img.size

            overlay = img.copy()

            label_fullname = os.path.join(self.args.lbl_dir, basename.replace(self.args.img_ext, '.xml'))
            if not os.path.exists(label_fullname):
                logger.warning('label file not found: %s', label_fullname)
                continue

            parser = helper.LabelXml(class_names, label_fullname)
            for class_id, cxywh in parser.get_boxes((ih, iw)):
                axywh = helper.cxywh2axywh(cxywh)
                x, y, w, h = axywh

                class_name = class_names[int(class_id)]
                # print(class_name, x, y, w, h)

                rect_color = (0, 255, 0)
                ImageDraw.Draw(overlay).rectangle([(x, y), (x+w, y+h)], outline=rect_color)

                fill_color = (255, 0, 0)
                stroke_color = (255, 255, 255)
                text_x, text_y = x, y
                ImageDraw.Draw(img).text((text_x, text_y), class_name, fill=fill_color, stroke_width=2, stroke_fill=stroke_color, font=font)
                ImageDraw.Draw(overlay).text((text_x, text_y), class_name, fill=fill_color, stroke_width=2, stroke_fill=stroke_color, font=font)

            alpha = 0.7
            overlay.putalpha(int(alpha * 255))
            img.alpha_composite(overlay)

            if args.preview:
                plt.imshow(img)
                cv2.waitKey(0)
                cv2.destroyAllWindows()
            else:
                out_fullname = os.path.join(self.args.out_dir, basename)
                # cv2.imwrite(out_fullname, img)
                img.convert('RGB').save(out_fullname)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--img-dir', type=str, default=r'C:\my\ws\static\fair1m1.0\train-part1\images')
    parser.add_argument('--img-ext', type=str, default='.tif')
    parser.add_argument('--lbl-dir', type=str, default=r'C:\my\ws\static\fair1m1.0\train-part1\labelXml')
    parser.add_argument('--out-dir', type=str, default=r'C:\my\ws\static\fair1m1.0\train-part1\images-box')
    parser.add_argument('--preview', action='store_true')
    args = parser.parse_args()
    Main(args).start()
